import os
import cv2
import shutil
from glob import glob
from tqdm import tqdm
from pascal_voc_writer import Writer


from utils.label_utils import parse_xml


def copy_to_dst():
    fd = r'F:\datasets\wb\seg_images\images'
    dst_fd = r'F:\datasets\wb\images'
    xmls = glob(fd+'/*.xml')
    for xml in tqdm(xmls):
        shutil.copy(xml, dst_fd)
        shutil.copy(xml.replace('.xml', '.jpg'), dst_fd)


def seg_to_pieces():
    fd = r'F:\datasets\wb\images'
    save_fd = r'F:\datasets\wb\carriages-seg'
    xmls = glob(fd + '/*.xml')
    xmls.sort()
    idx = 0
    with open('data.txt', 'w') as f:
        for xml in tqdm(xmls):
            data = parse_xml(xml)
            image = cv2.imread(xml.replace('.xml', '.jpg'))
            for i, box in enumerate(data.boxes):
                _, x1, y1, x2, y2 = box.box
                t_im = image[y1:y2, x1:x2]
                save_name = os.path.join(save_fd, str(idx).rjust(5,"0") + '-' + os.path.basename(xml).replace('.xml', '-' + str(i) + '.jpg'))
                f.write(save_name + ' \n')
                cv2.imwrite(save_name, t_im)
                idx = idx + 1


def main():
    fd = r'F:\datasets\wb\images'
    save_fd = r'F:\datasets\wb\carriages-seg'
    xmls = glob(fd + '/*.xml')
    xmls.sort()
    idx = 0
    with open('data.txt', 'r') as f:
        lines = f.readlines()
    dct = dict()
    for line in lines:
        s = line.strip().split(' ')
        if len(s) != 2:
            print(line)
        else:
            dct[s[0]] = s[1]
    pass
    for xml in tqdm(xmls):
        data = parse_xml(xml)
        res = []
        for i, box in enumerate(data.boxes):
            x1, y1, x2, y2 = box.box
            save_name = save_fd + '\\' + str(idx).rjust(5, "0") + '-' + os.path.basename(xml).replace('.xml', '-' + str(i) + '.jpg')
            if save_name not in dct.keys():
                print(xml)
            item = dct[save_name]
            res.append((x1, y1, x2, y2, item))
            idx = idx + 1
        if len(res) > 0:
            writer = Writer(xml.replace('.xml', '.jpg'), data.width, data.height)
            for item in res:
                x1, y1, x2, y2, c = item
                writer.addObject(c, x1, y1, x2, y2)
            writer.save(xml)


if __name__ == '__main__':
    main()
